{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9c74897e",
   "metadata": {},
   "source": [
    "## Evaluating a Health Assistant Agent Built with LangGraph\n",
    "\n",
    "In this notebook you will learn: \n",
    "\n",
    "- Evalauate the agent using [TaskCompletion metric](https://deepeval.com/docs/metrics-task-completion)\n",
    "- Change the hyperparameter to improve the agent's performance\n",
    "- Evaluate the agent again\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42aacfa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -U langgraph langchain langchain-community langchain-openai chromadb --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19b24062",
   "metadata": {},
   "source": [
    "Export you OPENAI_API_KEY as an environment variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3ae50209",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"<your-api-key>\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66b83cb5",
   "metadata": {},
   "source": [
    "### Health assistant agent built with LangGraph\n",
    "\n",
    "Given a user query, the agent will decide the best way to process the query. Here is the diagram of the agent:\n",
    " \n",
    "<img src=\"static/output.png\" alt=\"Agent Diagram\" height=\"300\" style=\"display: block; margin: 0 auto;\">\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d7e6786",
   "metadata": {},
   "source": [
    "We are keeping the model as `gpt-4o-mini` for the first iteration. Later in the same notebook we will evaluate the agent with `gpt-4` to see the performance difference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "97c268ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05b9af81",
   "metadata": {},
   "source": [
    "Pull the `manual.txt` which will form knowlege base of the agent "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b30092c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "!curl -o manual.txt \"https://confident-bucket.s3.us-east-1.amazonaws.com/manual.txt\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ff3686c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "from typing import Annotated, List, TypedDict, Literal\n",
    "\n",
    "from langchain_core.messages import BaseMessage, HumanMessage, AIMessage\n",
    "from langchain_core.tools import tool\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_community.vectorstores import Chroma\n",
    "from langchain_community.document_loaders import TextLoader\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.graph.message import add_messages\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "# Set API keys\n",
    "os.environ[\"OPENAI_API_KEY\"] = os.getenv(\"OPENAI_API_KEY\")\n",
    "\n",
    "\n",
    "class AgentState(TypedDict):\n",
    "    \"\"\"State schema for the RAG agent\"\"\"\n",
    "\n",
    "    messages: Annotated[List[BaseMessage], add_messages]\n",
    "    query: str\n",
    "    selected_tools: List[str]\n",
    "    retrieved_context: str\n",
    "    tool_outputs: List[str]\n",
    "    next_action: str\n",
    "\n",
    "\n",
    "# Initialize vector store with your knowledge base\n",
    "def setup_vector_store():\n",
    "    \"\"\"Set up your vector database with documents from local text file\"\"\"\n",
    "    # Load your documents from local text file\n",
    "    text_file_path = \"manual.txt\"  # Replace with your actual file path\n",
    "\n",
    "    try:\n",
    "        # Load the text file\n",
    "        loader = TextLoader(text_file_path, encoding=\"utf-8\")\n",
    "        docs = loader.load()\n",
    "\n",
    "        # Split documents\n",
    "        text_splitter = RecursiveCharacterTextSplitter(\n",
    "            chunk_size=500, chunk_overlap=50\n",
    "        )\n",
    "        doc_splits = text_splitter.split_documents(docs)\n",
    "        # Create vector store\n",
    "        embeddings = OpenAIEmbeddings()\n",
    "        vector_store = Chroma.from_documents(doc_splits, embeddings)\n",
    "\n",
    "        return vector_store.as_retriever()\n",
    "\n",
    "    except FileNotFoundError:\n",
    "        print(\n",
    "            f\"Error: File '{text_file_path}' not found. Please check the file path.\"\n",
    "        )\n",
    "        return None\n",
    "    except Exception as e:\n",
    "        print(f\"Error loading document: {str(e)}\")\n",
    "        return None\n",
    "\n",
    "\n",
    "retriever = setup_vector_store()\n",
    "\n",
    "\n",
    "######## TOOLS ########\n",
    "@tool\n",
    "def get_last_day_steps():\n",
    "    \"\"\"\n",
    "    Get the last day's steps from the database\n",
    "    \"\"\"\n",
    "    return random.randint(1000, 5000)\n",
    "\n",
    "\n",
    "@tool\n",
    "def get_last_day_average_heart_rate():\n",
    "    \"\"\"\n",
    "    Get the last day's average heart rate from the database\n",
    "    \"\"\"\n",
    "    return random.randint(60, 100)\n",
    "\n",
    "\n",
    "@tool\n",
    "def get_last_day_average_sleep_duration_in_hours():\n",
    "    \"\"\"\n",
    "    Get the last day's average sleep duration from the database\n",
    "    \"\"\"\n",
    "    return random.randint(3, 10)\n",
    "\n",
    "\n",
    "# Tool registry for dynamic selection\n",
    "tools = [\n",
    "    get_last_day_steps,\n",
    "    get_last_day_average_heart_rate,\n",
    "    get_last_day_average_sleep_duration_in_hours,\n",
    "]\n",
    "tool_registry = {tool.name: tool for tool in tools}\n",
    "\n",
    "\n",
    "######## TYPE DEFINITIONS ########\n",
    "class RouteQuery(BaseModel):\n",
    "    \"\"\"Schema for routing decisions\"\"\"\n",
    "\n",
    "    reasoning: str = Field(description=\"Reasoning for the routing decision\")\n",
    "    route: Literal[\"retrieval\", \"tools\", \"direct\"] = Field(\n",
    "        description=\"Where to route the query\"\n",
    "    )\n",
    "    tools_needed: List[str] = Field(\n",
    "        description=\"List of tools needed if route is 'tools'\"\n",
    "    )\n",
    "    retrieval_query: str = Field(\n",
    "        description=\"Optimized query for retrieval if route is 'retrieval'\"\n",
    "    )\n",
    "\n",
    "\n",
    "def router_node(state: AgentState) -> AgentState:\n",
    "    \"\"\"Route the query to appropriate processing path\"\"\"\n",
    "\n",
    "    system_prompt = \"\"\"You are an intelligent router that decides how to process user queries.\n",
    "    \n",
    "    Available options:\n",
    "    - 'retrieval': Query needs information from the knowledge base\n",
    "    - 'tools': Query needs external tools (web search, calculations, etc.)  \n",
    "    - 'direct': Query can be answered directly with general knowledge\n",
    "    \n",
    "    Available tools: {tools}\n",
    "    \n",
    "    Analyze the user query and decide the best routing approach. If tools are needed,\n",
    "    specify which ones. If retrieval is needed, optimize the query for better results.\"\"\"\n",
    "\n",
    "    user_query = state[\"messages\"][-1].content\n",
    "\n",
    "    structured_llm = llm.with_structured_output(\n",
    "        RouteQuery, method=\"function_calling\"\n",
    "    )\n",
    "\n",
    "    response = structured_llm.invoke(\n",
    "        [\n",
    "            {\n",
    "                \"role\": \"system\",\n",
    "                \"content\": system_prompt.format(tools=[t.name for t in tools]),\n",
    "            },\n",
    "            {\"role\": \"user\", \"content\": user_query},\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    return {\n",
    "        \"query\": user_query,\n",
    "        \"next_action\": response.route,\n",
    "        \"selected_tools\": response.tools_needed,\n",
    "        \"retrieved_context\": (\n",
    "            response.retrieval_query if response.route == \"retrieval\" else \"\"\n",
    "        ),\n",
    "    }\n",
    "\n",
    "\n",
    "def tool_execution_node(state: AgentState) -> AgentState:\n",
    "    \"\"\"Execute selected tools\"\"\"\n",
    "    tool_outputs = []\n",
    "    for tool_name in state[\"selected_tools\"]:\n",
    "        if tool_name in tool_registry:\n",
    "            tool = tool_registry[tool_name]\n",
    "            try:\n",
    "                # Use the original query for tool execution\n",
    "                output = tool.invoke({\"query\": state[\"query\"]})\n",
    "                tool_outputs.append(f\"{tool_name}: {output}\")\n",
    "            except Exception as e:\n",
    "                tool_outputs.append(f\"{tool_name}: Error - {str(e)}\")\n",
    "\n",
    "    return {\"tool_outputs\": tool_outputs}\n",
    "\n",
    "\n",
    "def retrieval_node(state: AgentState) -> AgentState:\n",
    "    \"\"\"Execute retrieval from vector database\"\"\"\n",
    "    query = state[\"retrieved_context\"] or state[\"query\"]\n",
    "\n",
    "    try:\n",
    "        # Retrieve relevant documents\n",
    "        docs = retriever.invoke(query)\n",
    "        context = \"\\n\\n\".join([doc.page_content for doc in docs])\n",
    "        return {\"retrieved_context\": context}\n",
    "    except Exception as e:\n",
    "        return {\"retrieved_context\": f\"Retrieval error: {str(e)}\"}\n",
    "\n",
    "\n",
    "def response_synthesis_node(state: AgentState) -> AgentState:\n",
    "    \"\"\"Synthesize final response from all available information\"\"\"\n",
    "\n",
    "    # Prepare context from various sources\n",
    "    context_parts = []\n",
    "\n",
    "    if state.get(\"retrieved_context\"):\n",
    "        context_parts.append(\n",
    "            f\"Knowledge Base Context:\\n{state['retrieved_context']}\"\n",
    "        )\n",
    "\n",
    "    if state.get(\"tool_outputs\"):\n",
    "        tool_context = \"\\n\".join(state[\"tool_outputs\"])\n",
    "        context_parts.append(f\"Tool Outputs:\\n{tool_context}\")\n",
    "\n",
    "    context = \"\\n\\n\".join(context_parts)\n",
    "\n",
    "    system_prompt = \"\"\"You are a helpful assistant that synthesizes information from multiple sources.\n",
    "    \n",
    "    Use the provided context to answer the user's question accurately and comprehensively.\n",
    "    If using information from the context, be sure to reference it appropriately.\n",
    "    If the context doesn't contain enough information, acknowledge this limitation.\"\"\"\n",
    "\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": f\"Question: {state['query']}\\n\\nContext:\\n{context}\",\n",
    "        },\n",
    "    ]\n",
    "\n",
    "    response = llm.invoke(messages)\n",
    "\n",
    "    return {\"messages\": [AIMessage(content=response.content)]}\n",
    "\n",
    "\n",
    "def intial_route_decision(state: AgentState) -> str:\n",
    "    \"\"\"Determine next node based on routing decision\"\"\"\n",
    "    next_action = state.get(\"next_action\", \"direct\")\n",
    "\n",
    "    if next_action == \"tools\":\n",
    "        return \"tools\"\n",
    "\n",
    "    if next_action == \"retrieval\":\n",
    "        return \"retrieval\"\n",
    "\n",
    "    return \"retrieval\"\n",
    "\n",
    "\n",
    "def create_rag_graph():\n",
    "    \"\"\"Create and compile the RAG workflow graph\"\"\"\n",
    "\n",
    "    # Initialize the state graph\n",
    "    workflow = StateGraph(AgentState)\n",
    "\n",
    "    # Add nodes\n",
    "    workflow.add_node(\"router\", router_node)\n",
    "    workflow.add_node(\"retrieval\", retrieval_node)\n",
    "    workflow.add_node(\"tools\", tool_execution_node)\n",
    "    workflow.add_node(\"synthesis\", response_synthesis_node)\n",
    "\n",
    "    # define edges\n",
    "    workflow.add_edge(START, \"router\")\n",
    "    workflow.add_conditional_edges(\n",
    "        \"router\",\n",
    "        intial_route_decision,\n",
    "        {\n",
    "            \"retrieval\": \"retrieval\",\n",
    "            \"tools\": \"tools\",\n",
    "            \"synthesis\": \"synthesis\",\n",
    "        },\n",
    "    )\n",
    "\n",
    "    workflow.add_edge(\"retrieval\", \"synthesis\")\n",
    "    workflow.add_edge(\"tools\", \"synthesis\")\n",
    "    workflow.add_edge(\"synthesis\", END)\n",
    "\n",
    "    return workflow.compile()\n",
    "\n",
    "\n",
    "# Create the graph\n",
    "app = create_rag_graph()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c32deb1",
   "metadata": {},
   "source": [
    "Now we have the graph, we can run the agent with the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e7aa474",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_state = {\n",
    "    \"query\": \"\",\n",
    "    \"selected_tools\": [],\n",
    "    \"retrieved_context\": \"\",\n",
    "    \"tool_outputs\": [],\n",
    "    \"next_action\": \"\",\n",
    "}\n",
    "\n",
    "\n",
    "def run_rag_query(query: str):\n",
    "    \"\"\"Run a query through the RAG system\"\"\"\n",
    "\n",
    "    initial_state[\"messages\"] = [HumanMessage(content=query)]\n",
    "    result = app.invoke(initial_state)\n",
    "    final_message = result[\"messages\"][-1]\n",
    "    return final_message.content\n",
    "\n",
    "\n",
    "run_rag_query(\"What is the average heart rate of the user?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6084b818",
   "metadata": {},
   "source": [
    "### Evaluate the agent\n",
    "\n",
    "[DeepEval](https://deepeval.com/) provides a `CallbackHandler` for LangGraph and LangChain agents to evaluate (and trace) the agents. \n",
    "\n",
    "\n",
    "> (Pro Tip) View your Agent's trace and publish test runs on [Confident AI](https://www.confident-ai.com/). Apart from this you get an in-house dataset editor and more advaced tools to monitor and enventually improve your Agent's performance. Get your API key from [here](https://app.confident-ai.com/)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a066c341",
   "metadata": {},
   "source": [
    "OPTIONAL: Set CONFIDENT_API_KEY as an environment variable to publish test results on Confident AI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dce1b09",
   "metadata": {},
   "outputs": [],
   "source": [
    "!export CONFIDENT_API_KEY=your-api-key"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "560fe4bf",
   "metadata": {},
   "source": [
    "Initialize the CallbackHandler and pass TaskCompletionMetric to it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23a84a55",
   "metadata": {},
   "outputs": [],
   "source": [
    "from deepeval.integrations.langchain import CallbackHandler\n",
    "from deepeval.metrics import TaskCompletionMetric\n",
    "\n",
    "\n",
    "def run_rag_query(query: str):\n",
    "    \"\"\"Run a query through the RAG system\"\"\"\n",
    "\n",
    "    initial_state[\"messages\"] = [HumanMessage(content=query)]\n",
    "\n",
    "    result = app.invoke(\n",
    "        initial_state,\n",
    "        config={\n",
    "            \"callbacks\": [\n",
    "                CallbackHandler(\n",
    "                    metrics=[\n",
    "                        TaskCompletionMetric(strict_mode=True, async_mode=False)\n",
    "                    ]\n",
    "                )\n",
    "            ]  # pass the metrics to the callback handler\n",
    "        },\n",
    "    )\n",
    "\n",
    "    final_message = result[\"messages\"][-1]\n",
    "    return final_message.content"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6aa79b3d",
   "metadata": {},
   "source": [
    "### Pull the dataset\n",
    "For tutorial purposes, we will use the public dataset of health queries. You can use your own dataset as well. Refer to the [docs](https://deepeval.com/docs/evaluation-end-to-end-llm-evals#setup-your-test-environment) to learn more about how to create your own dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b84aa705",
   "metadata": {},
   "outputs": [],
   "source": [
    "from deepeval.dataset import EvaluationDataset\n",
    "\n",
    "dataset = EvaluationDataset()\n",
    "dataset.pull(alias=\"health_rag_queries\", public=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73080eb9",
   "metadata": {},
   "source": [
    "Run evals using dataset iterator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a9b7871",
   "metadata": {},
   "outputs": [],
   "source": [
    "for golden in dataset.evals_iterator():\n",
    "    run_rag_query(golden.input)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cb2c55d",
   "metadata": {},
   "source": [
    "### Change the model to gpt-4 and evaluate again\n",
    "\n",
    "Now we will change the model to `gpt-4`, redefine the nodes and evaluate the agent again.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d27fefa6",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = ChatOpenAI(model=\"gpt-4\", temperature=0)\n",
    "app = create_rag_graph()\n",
    "\n",
    "\n",
    "def run_rag_query(query: str):\n",
    "\n",
    "    initial_state[\"messages\"] = [HumanMessage(content=query)]\n",
    "    result = app.invoke(\n",
    "        initial_state,\n",
    "        config={\n",
    "            \"callbacks\": [\n",
    "                CallbackHandler(\n",
    "                    metrics=[\n",
    "                        TaskCompletionMetric(strict_mode=True, async_mode=False)\n",
    "                    ]\n",
    "                )\n",
    "            ]\n",
    "        },\n",
    "    )\n",
    "    final_message = result[\"messages\"][-1]\n",
    "    return final_message.content\n",
    "\n",
    "\n",
    "for golden in dataset.evals_iterator():\n",
    "    run_rag_query(golden.input)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60a507f6",
   "metadata": {},
   "source": [
    "Try changing other hyperparameters of the model and evaluate the agent again."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
